from src.agents.outline_writer import outlineWriter
from src.agents.writer import subsectionWriter
from src.database import PaperDatabase
'''
# 创建数据库实例
db = PaperDatabase(
    db_type="local",
    local_config={
        "db_path": "./database",
        "embedding_model": "./embedding_models/AI-ModelScope/nomic-embed-text-v1",
        "device": "cuda"  # 或 "cpu"
    }
)
'''

# 创建数据库实例
db = PaperDatabase(
    db_type="atomgit",
    atomgit_config={
        "base_url": "http://180.184.65.98:38880/atomgit"
    }
)

model='glm-4-plus'
api_key="99fbf79da7e241929139d3e9fcc0e3d2.M8PE9i2FJiVd3RBy"
api_url="https://open.bigmodel.cn/api/paas/v4/"
topic="有哪些方法可以提升大模型的规划能力，各自优劣是什么？"


outline_writer = outlineWriter(
    model=model, 
    api_key=api_key, 
    api_url=api_url, 
    database=db
)
subsection_writer = subsectionWriter(
    model=model, 
    api_key=api_key, 
    api_url=api_url, 
    database=db
    )

outline = outline_writer.draft_outline(topic)
# print(outline)


raw_survey, raw_survey_with_references, raw_references, refined_survey, refined_survey_with_references, refined_references = subsection_writer.write(topic=topic, outline=outline)

with open(f'./output/{topic}.md', 'w') as f:
    f.write(refined_survey_with_references)
